from typing import List, Literal, Optional, Dict, Any
from langchain_community.embeddings import JinaEmbeddings, ZhipuAIEmbeddings
from MyEmbeddingModel import CustomAPIEmbeddings,InstructorEmbeddings

class EmbeddingClient:
    def __init__(
        self,
        mode: Literal["jina", "zhipu", "myapi"],
        model_name: Optional[str] = None,
        api_key: Optional[str] = None,
        model_kwargs: Optional[Dict[str, Any]] = None,
    ):
        """
        初始化嵌入客户端
        
        Args:
            mode: 嵌入模型类型 ("jina", "zhipu", "myapi")
            model_name: 模型名称/ID
            api_key: 模型API密钥（可为None，如果model_kwargs中已提供）
            model_kwargs: 模型特定的参数字典
        """
        self.mode = mode
        self.model_name = model_name
        self.model_kwargs = model_kwargs or {}
        self.embedding_client = None
        self.embedding_model = None

        # 如果api_key参数提供了但model_kwargs中没有，则添加到model_kwargs
        if api_key is not None and "api_key" not in self.model_kwargs:
            self.model_kwargs["api_key"] = api_key
            
        # 初始化对应的嵌入模型
        if self.mode == "jina":
            self._init_jina()
        elif self.mode == "zhipu":
            self._init_zhipu()
        elif self.mode == "myapi":
            self._init_myapi()
        else:
            raise ValueError(f"Unsupported mode: {self.mode}")

    def _init_jina(self):
        """初始化Jina嵌入模型"""
        if "api_key" not in self.model_kwargs:
            raise ValueError("Jina embeddings require 'jina_api_key' in model_kwargs")
        
        self.embedding_client = JinaEmbeddings(
            api_key=self.model_kwargs["api_key"],
            model_name=self.model_name or "jina-embeddings-v2-base-en",
            **{k: v for k, v in self.model_kwargs.items() if k != "api_key"}
        )

    def _init_zhipu(self):
        """初始化ZhipuAI嵌入模型"""
        if "api_key" not in self.model_kwargs:
            raise ValueError("ZhipuAI embeddings require 'zhipuai_api_key' in model_kwargs")
        
        self.embedding_client = ZhipuAIEmbeddings(
            api_key=self.model_kwargs["api_key"],
            model=self.model_name or "embedding-2",
            **{k: v for k, v in self.model_kwargs.items() if k != "api_key"}
        )

    def _init_myapi(self):
        """初始化自定义API嵌入模型"""
        # 自定义模型可能需要api_key或其他参数，都在model_kwargs中传递
        self.embedding_client = CustomAPIEmbeddings(
            api_key=self.model_kwargs["api_key"],
            **{k: v for k, v in self.model_kwargs.items() if k != "api_key"}
        )

    def get_embedding_model(self):
        return InstructorEmbeddings(embedding_client = self.embedding_client)

    def embed(self, text: str) -> List[float]:
        """
        生成文本的嵌入向量
        
        Args:
            text: 要嵌入的文本
            
        Returns:
            嵌入向量列表
        """
        try:
            # 优先尝试embed_documents方法（批量处理）
            if hasattr(self.embedding_model, 'embed_documents'):
                return self.embedding_model.embed_documents([text])[0]
            # 其次尝试embed_query方法（单条处理）
            elif hasattr(self.embedding_model, 'embed_query'):
                return self.embedding_model.embed_query(text)
            else:
                raise AttributeError("Embedding model doesn't have required embedding methods")
        except Exception as e:
            raise RuntimeError(f"Embedding failed: {str(e)}") from e